Problem

group_norm_backward picks the dtype its dW/dB gradient accumulators are cast to before the atomic_add into the DW/DB buffers:

triton_dtype = tl.float32 if X.dtype == torch.float32 else tl.bfloat16

The DW/DB buffers are allocated with W.dtype / B.dtype. For an fp16 model that's an fp16 buffer, but triton_dtype is bfloat16 — so the fp16 gradients are rounded to bfloat16 (8 mantissa bits) and then atomic-added into an fp16 buffer (10 mantissa bits), losing precision for no reason. bf16 and fp16 are not interchangeable.

Fix

Map fp16 -> tl.float16 so the accumulator/atomic dtype matches the buffer dtype. bf16 inputs are unchanged.

Testing done (NVIDIA T4)

GroupNorm forward+backward vs torch.nn.functional.group_norm, comparing the weight gradient dW[0] (ref = 4.981):

                 fp32            fp16
before (main)    dW=4.981        dW=4.996   (bf16-rounded accumulator)
after  (this PR) dW=4.981        dW=4.984   (closer to the fp32 reference)

Both pass the tolerance check; the fix moves the fp16 gradient measurably closer to the reference. Honesty note: on the Triton version tested the old bf16 -> fp16 atomic_add is silently coerced rather than erroring, so this is a precision/type-consistency fix, not a crash fix.